package com.learn.unionset;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

public class UnionSetDemo {
    public static class Node<V> {
        private V value;

        public Node(V value) {
            this.value = value;
        }
    }

    public static class UnionSet<V> {
        public HashMap<V, Node<V>> nodes;
        public HashMap<Node<V>, Node<V>> parentMaps;
        public HashMap<Node<V>, Integer> sizeMap;

        public UnionSet(List<V> values) {
            for (V v : values) {
                Node node = new Node(v);
                nodes.put(v, node);
                parentMaps.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        public Node<V> findParent(Node<V> node) {
            Stack<Node> stack = new Stack<>();
            while (node != parentMaps.get(node)) {
                stack.push(node);
                node = parentMaps.get(node);
            }
            while (!stack.isEmpty()) {
                parentMaps.put(stack.pop(), node);
            }

            return node;
        }


        public boolean isSameSet(V a, V b) {
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return false;
            }
            //说明nodes初始化的时候，已经存在，则需要判断它们的父类是不是同一个，是则在同一个集合

            return findParent(nodes.get(a)) == findParent((nodes.get(b)));

        }

        public void union(V a, V b) {
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return;
            }
            Node<V> aHead = findParent(nodes.get(a));
            Node<V> bHead = findParent(nodes.get(b));

            //它们的头部不相等，代表不属于同一集合
            if (aHead != bHead) {
                //bHead需要指向aHead
                if (sizeMap.get(aHead) > sizeMap.get(bHead)) {
                    parentMaps.put(bHead, aHead);
                    sizeMap.put(aHead, sizeMap.get(aHead) + sizeMap.get(bHead));
                    sizeMap.remove(bHead);
                } else {
                    parentMaps.put(aHead, bHead);
                    sizeMap.put(bHead, sizeMap.get(aHead) + sizeMap.get(bHead));
                    sizeMap.remove(aHead);
                }

            }

        }
    }
}
